import cv2
import numpy as np
import glob
from pprint import pprint
image_paths = glob.glob("img/*.png")


# 内参 : 描述相机 内部的光学和几何特性，将 3D相机坐标系中的点 映射到 2D图像像素坐标系。
# 外惨 : 描述相机 在世界坐标系中的位姿（旋转和平移），用于将 世界坐标 转换到 相机坐标
pattern_size = (9, 11)

# 世界坐标系下的 3D 点（z = 0）
objp = np.zeros((np.prod(pattern_size), 3), np.float32)
objp[:, :2] = np.mgrid[0 : pattern_size[0], 0 : pattern_size[1]].T.reshape(-1, 2)

objpoints = []  # 3D点
imgpoints = []  # 2D点

for fname in image_paths:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 找到棋盘格角点
    ret, corners = cv2.findChessboardCorners(
        gray, pattern_size, cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
    )

    if ret:
        objpoints.append(objp)
        # 亚像素精确角点
        corners_subpix = cv2.cornerSubPix(
            gray,
            corners,
            (11, 11),
            (-1, -1),
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001),
        )
        imgpoints.append(corners_subpix)

# -------------------------------
# 3. 角点匹配（OpenCV 内部自动处理）
# -------------------------------
# 匹配由角点顺序自动对应，不需手动匹配

# -------------------------------
# 4. 构建投影变换 +
# 5. 最小二乘法估计内外参
# -------------------------------
# 相机标定（使用最小二乘优化）
ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
    objpoints, imgpoints, gray.shape[::-1], None, None
)

# -------------------------------
# 6. 分析内参与外参
# -------------------------------

print("相机内参矩阵 (Camera Matrix):")
print(camera_matrix)
print("\n畸变系数 (Distortion Coefficients):")
print(dist_coeffs.ravel())

for i, (rvec, tvec) in enumerate(zip(rvecs, tvecs)):
    R, _ = cv2.Rodrigues(rvec)
    print(f"\n图像 {i + 1} 外参 (外部矩阵):")
    print("旋转矩阵 R:")
    print(R)
    print("平移向量 t:")
    print(tvec.ravel())

# 可选：重投影误差分析
total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(
        objpoints[i], rvecs[i], tvecs[i], camera_matrix, dist_coeffs
    )
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
    total_error += error
print(f"\n平均重投影误差: {total_error / len(objpoints)}")


